﻿# _*_coding:utf-8_*_

# 导入依赖包，初始化一些常量
import tensorflow as tf
import numpy as np
import reader
import collections

"""
import collections
TRAIN_DATA = "./data/ptb.train.txt"  # 训练数据路径
TEST_DATA = "./data/ptb.test.txt"  # 测试数据路径
EVAL_DATA = "./data/ptb.valid.txt"  # 验证数据路径

"""

DATA_PATH = "data"
init_scale = 0.1  # 相关参数的初始值为随机均匀分布，范围是[-init_scale,+init_scale]
LEARNING_RATE = 1.0  # 学习速率，此值还会在模型学习过程中下降
MAX_GRAD_NORM = 5  # 控制梯度膨胀的参数，如果梯度向量的L2模超过max_grad_norm，则等比例缩小

NUM_LAYERS = 2  # LSTM层数

HIDDEN_SIZE = 200  # 隐层单元数目，每个词会表示成[hidden_size]大小的向量
max_epoch = 4  # epoch<max_epoch时，lr_decay值=1,epoch>max_epoch时,lr_decay逐渐减小
max_max_epoch = 13  # 完整的文本要循环的次数
KEEP_PROB = 0.5  # 不dropout的概率,1.0为不丢弃
lr_decay = 0.5  # 学习速率衰减指数

VOCAB_SIZE = 10000  # 单词数量(这份训练数据中单词刚好10000种)

# NUM_STEPS=35        # 分隔句子的粒度大小，每次会把num_steps个单词划分为一句话(但是本模型与seq2seq模型不同，它仅仅是1对1模式，句子长度应该没有什么用处)。
# BATCH_SIZE = 20     # 和num_steps共同作用，要把原始训练数据划分为batch_size组，每组划分为n个长度为num_steps的句子。

TRAIN_BATCH_SIZE = 20  # 训练数据batch的大小,要把原始训练数据划分为TRAIN_BATCH_SIZE组，每组划分为n个长度为TRAIN_NUM_STEP的句子。
TRAIN_NUM_STEP = 35  # 训练数据的截断长度，也可以看作是序列的长度,分隔句子的粒度大小，每次会把TRAIN_NUM_STEP个单词划分为一句话

EVAL_BATCH_SIZE = 1
EVAL_NUM_STEP = 1

NUM_EPOCH = 2  # 使用训练数据的轮数
SHARE_EMB_AND_SOFTMAX = True  # 在softmax层和词向量层之间共享参数

# 词 向量化 转成向量
def read_data(file_path):
    """
    读取文件数据，将文本中的词转换成词空间中对应的索引
    :param file_path: 文件路径
    :return: 由数值取代后的文本词列表
    """

    # 采用TensorFlow中的读取文件的方法去读取文件
    with tf.gfile.GFile(file_path, "r") as f:
        # 将文本读取出来，并且进行分词，将换行符替换成<eos>,eos的意思就是end of sentence
        word_list = f.read().replace("\n", "<eos>").split()
        # 对分词后的列表进行统计，统计每个单词出现的数量, 返回的数据类型Counter({'jiang': 2, 'zhang': 1})
        counter = collections.Counter(word_list)
        # 对每个词按照词频排序，对于词频相同的按词本身进行排序，返回的数据类型[('jiang', 2), ('zhang', 1)]
        count_pairs = sorted(counter.items(), key=lambda x: (-x[1], x[0]))
        # 取出单词元组，返回的数据类型words=('jiang', 'zhang')
        words, _ = list(zip(*count_pairs))
        # 将上面排序后的词（无重复的词空间）标记索引数值，返回的数据类型{'jiang': 0, 'zhang': 1}
        word_to_id = dict(zip(words, range(len(words))))
        # 将文本中所有的单词用索引数值取代，组成新的列表
        id_list = [word_to_id[word] for word in word_list if word in word_to_id]

    return id_list


def make_batches(id_list, batch_size, num_steps):
    """
    将原始的数据集转换成mini-batch进行训练
    :param id_list: 原始的数值文本列表
    :param batch_size:
    :param num_steps: 一个样本的序列长度
    :return: 整个样本转换后的batchs数据
    """
    # 计算总的batch数量。每个batch包含的单词数量是batch_size * num_steps，batch_size为一个batch中样本的数量，
    # num_steps为一个样本的序列长度，因此num_batchs表示整个训练文本能分成的batch的数量
    num_batches = (len(id_list) - 1) // (batch_size * num_steps)

    # 根据上面分配好的num_batchs, batch_size, num_steps参数来构建数据集，先取出能整除的序列长度
    data = np.array(id_list[: num_batches * batch_size * num_steps])
    # reshape取出来的序列（一维数组）成二维数组，因为是序列数据，所以行为batch_size，列为num_batchs * num_steps，之后训练在横向上分割
    data = np.reshape(data, [batch_size, num_batches * num_steps])
    # 将数据在axis=1的轴上分割，分割的数量就是num_batchs
    data_batches = np.split(data, num_batches, axis=1)

    # 因为是根据前面的词预测后面的词，因此输出的值要往后移一位，其余操作和上面的一致
    label = np.array(id_list[1: num_batches * batch_size * num_steps + 1])
    label = np.reshape(label, [batch_size, num_batches * num_steps])
    label_batches = np.split(label, num_batches, axis=1)

    return list(zip(data_batches, label_batches))


class PTBModel(object):
    def __init__(self, is_training, batch_size, num_steps):
        # 记录使用的batch大小和截断长度（也就是一个样本的序列长度）
        self.batch_size = batch_size
        self.num_steps = num_steps

        # 初始化输入数据的维度，定义每一步的输入和预期输出。两者的维度都是[batch_size, num_steps]，batch_size时间上就是指一个batch中样本的数量
        self.input_data = tf.placeholder(tf.int32, [batch_size, num_steps])
        self.targets = tf.placeholder(tf.int32, [batch_size, num_steps])

        # 多层LSTM，设置dropout,定义lstm cell的结构，dropout相当于装饰器直接包裹在lstm_cell上,此时的cell是垂直方向的，所以for循环中的值是NUM_LAYERS）
        # lstm_cell = tf.nn.rnn_cell.BasicLSTMCell( HIDDEN_SIZE )
        lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(HIDDEN_SIZE, forget_bias=1.0, state_is_tuple=True)
        if is_training:
            lstm_cell = tf.nn.rnn_cell.DropoutWrapper(
                lstm_cell, output_keep_prob=KEEP_PROB)
        cell = tf.nn.rnn_cell.MultiRNNCell([lstm_cell] * NUM_LAYERS, state_is_tuple=True)

        # 初始化最初的状态值，即全为0的向量。这个量只在每个epoch初始化第一个batch时使用
        # 返回[batch_size, 2*len(cells)],或者[batch_size, s]，至于为什么是2 * ，这是因为初始值有两个h0和C0
        self.initial_state = cell.zero_state(batch_size, tf.float32)

        # 定义单词的词向量矩阵，这里用get_variable，之后在测试时就可以实现参数共享了
        # VOCAB_SIZE是指词向量空间中词的个数（在这里起始是len(words)的长度，也等于10000），HIDDEN_SIZE是值词嵌入之后的词向量长度
        embedding = tf.get_variable("embedding", [VOCAB_SIZE, HIDDEN_SIZE])

        # 将输入的单词转化为词向量，相当于将每个序列中的单词按照索引（之前转化为了数值在这里就很方便的)
        # 直接将序列中的每个词在已经训练好的词空间中寻找对应的向量，此空间应该是二维的，输出的结果应该是三维的，也就是batch_size * num_steps * HIDDEN_SIZE
        # 将id转换为词向量，从 batch_size * num_steps  到 batch_size * num_steps * HIDDEN_SIZE
        inputs = tf.nn.embedding_lookup(embedding, self.input_data)

        # 只在训练时使用dropout来训练词向量
        if is_training:
            inputs = tf.nn.dropout(inputs, KEEP_PROB)

        # 定义输出列表。在这里先将不同时刻LSTM结构的输出收集起来，再一起提供给softmax层，在这里是实现时间序上的cell输出
        outputs = []

        # 训练
        state = self.initial_state
        with tf.variable_scope("RNN"):
            for time_step in range(num_steps):
                if time_step > 0:
                    # 实现在同一个variable_scope下共享参数
                    tf.get_variable_scope().reuse_variables()
                # 从这里就可以看出之前embedding输出的是三维，我们根据时间序取出词来进行训练
                # cell_output shape=(batch_size, HIDDEN_SIZE)
                cell_output, state = cell(inputs[:, time_step, :], state)
                outputs.append(cell_output)
        # 把输出队列展开成[batch_size, num_steps * hidden_size]，然后再reshape成[batch_size * num_steps, hidden_size]
        output = tf.reshape(tf.concat(outputs, 1), [-1, HIDDEN_SIZE])#返回一个个单词的词向量

        # 一个全连接得到最后的结果
        # softmax层：将RNN在每个位置的输出转化为各个单词的logits
        # weight的shape是[HIDDEN_SIZE, VOCAB_SIZE]
        # weight = tf.get_variable( "weight", [ HIDDEN_SIZE, VOCAB_SIZE ] )
        if SHARE_EMB_AND_SOFTMAX:
            weight = tf.transpose(embedding)
        else:
            weight = tf.get_variable("weight", [HIDDEN_SIZE, VOCAB_SIZE])
        bias = tf.get_variable("bias", [VOCAB_SIZE])
        # 算出最终输出的logits，用于之后的交叉熵和softmax计算
        logits = tf.matmul(output, weight) + bias

        # 计算交叉熵作为loss,
        loss = tf.contrib.legacy_seq2seq.sequence_loss_by_example(
            [logits],#二维张量
            [tf.reshape(self.targets, [-1])],
            [tf.ones([batch_size * num_steps], dtype=tf.float32)]
        )

        """
        #定义交叉熵损失函数和平均损失函数，返回的loss是和labels、logits相同的shape
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=tf.reshape(self.targets, [-1]),
                                                              logits=logits)
        """

        # 求平均损失
        self.cost = tf.reduce_sum(loss) / batch_size
        # 该状态用来存储训练完一个batch之后的状态，在训练下一个batch时，会将该状态作为初始状态，这个在run_epoch函数中控制的
        self.final_state = state

        # 只在训练模型时定义反向传播操作
        if not is_training:
            return

        # tf.trainable_variables()返回所有训练的变量, 拿到所有的训练参数，用于之后的梯度下降更新参数
        trainable_variables = tf.trainable_variables()

        # 通过clip_by_global_norm避免膨胀梯度截断控制梯度大小，是为了避免梯度弥散和梯度爆炸，将梯度控制在某一范围内
        grads, _ = tf.clip_by_global_norm(
            # tf.gradients（）计算梯度 然后处理梯度
            tf.gradients(
                self.cost, trainable_variables
            ),
            MAX_GRAD_NORM
        )

        # 定义优化方法
        optimizer = tf.train.GradientDescentOptimizer(LEARNING_RATE)
        # 定义训练步骤
        self.train_op = optimizer.apply_gradients(
            zip(grads, trainable_variables)
        )


# 训练过程函数,用来控制模型的训练
def run_epoch(session, model, data, train_op, output_log):
    # 训练模型和模型预测。使用给定的模型model在数据data上运行train_op并返回在全部数据上的perplexity值。
    # 计算平均perplexity的辅助变量，perplexity表示在模型生成一句话时下一个词有perplexity个合理的选择，认为perplexity小于100都是比较好的
    # 结果，该值越小，说明模型越好，也可以认为模型预测的精确度越高

    # 存储损失值
    total_costs = 0.0
    iters = 0
    state = session.run(model.initial_state)

    # step = 0
    # [x,y] = reader.ptb_producer( data, model.batch_size, model.num_steps )
    # coord = tf.train.Coordinator()
    # tf.train.start_queue_runners(session, coord=coord)

    # 训练一个epoch
    for step, (x, y) in enumerate(
            reader.ptb_iterator(data, model.batch_size, model.num_steps)
    ):
        # [a,b] = session.run([x,y])
        # if a.size != model.batch_size * model.num_steps :
        #     break
        cost, state, _ = session.run(
            [model.cost, model.final_state, train_op],
            {
                model.input_data: x,
                model.targets: y,
                model.initial_state: state
            }
        )
        total_costs += cost
        iters += model.num_steps
        step += 1
        if output_log and step % 100 == 0:
            print("After %d steps, perplexity is %.3f" % (step, np.exp(total_costs / iters)))

            # pplx是语言模型perplexity指标
    pplx = np.exp(total_costs / iters)
    return step, pplx


# 定义main函数
def main(_):
    # with tf.Graph().as_default():

    train_data, valid_data, test_data, _ = reader.ptb_raw_data(DATA_PATH)
    # 定义初始化函数, 用于决定之后的variable_scope中的变量的初始值取值范围
    initializer = tf.random_uniform_initializer(-0.05, 0.05)

    # variable_scope为变量空间，当reuse=true时共享变量。
    with tf.variable_scope("language_model", reuse=None, initializer=initializer):
        train_model = PTBModel(True, TRAIN_BATCH_SIZE, TRAIN_NUM_STEP)

    with tf.variable_scope("language_model", reuse=True, initializer=initializer):
        eval_model = PTBModel(False, EVAL_BATCH_SIZE, EVAL_NUM_STEP)
    saver = tf.train.Saver()
    # 保存模型路徑
    model_save_path = r'F:\Resources\model\my_model\my_model'

    # 训练模型
    with tf.Session() as session:
        tf.global_variables_initializer().run()
        # 训练，每次训练后用valid数据测试
        for i in range(NUM_EPOCH):
            print("In iteration : %d " % (i + 1))
            run_epoch(session, train_model, train_data, train_model.train_op, True)

            valid_perplexity = run_epoch(session, eval_model, valid_data, tf.no_op(), False)[1]
            print("Epoch: %d Validation Perplexity : %.3f" % (i + 1, valid_perplexity))
        saver.save(session,save_path=model_save_path)
        # 調用模型
        saver.restore(session, model_save_path)
        # 在最终测试集上进行测试
        test_perplexity = run_epoch(session, eval_model, test_data, tf.no_op(), False)
        print("test Perplexity : %.3f" % (test_perplexity))


if __name__ == "__main__":
    tf.app.run()
# session = tf.Session()
# coord = tf.train.Coordinator()
# tf.train.start_queue_runners(session, coord=coord)
# [a,b] = session.run([x,y])
# print(a)
# [a,b] = session.run([x,y])
# print(a)
# # print("444")
